/*
	TRAINING SET LABEL FILE (train-labels-idx1-ubyte):
	[offset] [type]          [value]          [description] 
	0000     32 bit integer  0x00000801(2049) magic number (MSB first) 
	0004     32 bit integer  60000            number of items 
	0008     unsigned byte   ??               label 
	0009     unsigned byte   ??               label 
	........ 
	xxxx     unsigned byte   ??               label
	The labels values are 0 to 9.
	
	TRAINING SET IMAGE FILE (train-images-idx3-ubyte):
	[offset] [type]          [value]          [description] 
	0000     32 bit integer  0x00000803(2051) magic number 
	0004     32 bit integer  60000            number of images 
	0008     32 bit integer  28               number of rows 
	0012     32 bit integer  28               number of columns 
	0016     unsigned byte   ??               pixel 
	0017     unsigned byte   ??               pixel 
	........ 
	xxxx     unsigned byte   ??               pixel
	Pixels are organized row-wise. Pixel values are 0 to 255. 0 means background (white), 255 means foreground (black).

	TEST SET LABEL FILE (t10k-labels-idx1-ubyte):
	[offset] [type]          [value]          [description] 
	0000     32 bit integer  0x00000801(2049) magic number (MSB first) 
	0004     32 bit integer  10000            number of items 
	0008     unsigned byte   ??               label 
	0009     unsigned byte   ??               label 
	........ 
	xxxx     unsigned byte   ??               label
	The labels values are 0 to 9.
	
	TEST SET IMAGE FILE (t10k-images-idx3-ubyte):
	[offset] [type]          [value]          [description] 
	0000     32 bit integer  0x00000803(2051) magic number 
	0004     32 bit integer  10000            number of images 
	0008     32 bit integer  28               number of rows 
	0012     32 bit integer  28               number of columns 
	0016     unsigned byte   ??               pixel 
	0017     unsigned byte   ??               pixel 
	........ 
	xxxx     unsigned byte   ??               pixel
 */
package DataSet;

import java.io.DataInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.util.Date;

public class MnistDataSet {
	private String dataPath;
	private final static String train_labels_idx_file = "train-labels.idx1-ubyte";
	private final static String train_images_file = "train-images.idx3-ubyte";
	private final static String test_labels_idx_file = "t10k-labels.idx1-ubyte";
	private final static String test_images_file = "t10k-images.idx3-ubyte";

	private int[][] train_Images;
	private int[] train_Labels;

	private int[][] test_Images;
	private int[] test_Labels;

	private int labels_N = 10;

	private int rows, cols;
	private int train_N, test_N;

	public int getLabels_N() {
		return labels_N;
	}

	public int getRows() {
		return rows;
	}

	public int getCols() {
		return cols;
	}

	public int getImageDataLen() {
		return rows * cols;
	}

	public int getTrain_N() {
		return train_N;
	}

	public int getTest_N() {
		return test_N;
	}

	public int[][] getTrainImages() {
		return train_Images;
	}

	public double[][] getTrainDoubleImages() {
		double[][] res = new double[train_N][getImageDataLen()];
		for (int i = 0; i < train_N; i++) {
			for (int k = 0; k < getImageDataLen(); k++) {
				res[i][k] = train_Images[i][k];
			}
		}
		return res;
	}

	public double[][] getTestDoubleImages() {
		double[][] res = new double[test_N][getImageDataLen()];
		for (int i = 0; i < test_N; i++) {
			for (int k = 0; k < getImageDataLen(); k++) {
				res[i][k] = test_Images[i][k];
			}
		}
		return res;
	}

	public int[][] getTestImages() {
		return test_Images;
	}

	public int[] getTrainLabels() {
		return train_Labels;
	}

	public int[] getTestLabels() {
		return test_Labels;
	}

	public int[][] getCrossEntropyTrainLabels() {
		int[][] res = new int[train_N][labels_N];
		for (int i = 0; i < train_N; i++) {
			for (int k = 0; k < labels_N; k++) {
				if (k == train_Labels[i]) {
					res[i][k] = 1;
				} else {
					res[i][k] = 0;
				}
			}
		}
		return res;
	}

	public Integer[][] getCrossEntropyTestLabels() {
		Integer[][] res = new Integer[test_N][labels_N];
		for (int i = 0; i < test_N; i++) {
			for (int k = 0; k < labels_N; k++) {
				if (k == test_Labels[i]) {
					res[i][k] = 1;
				} else {
					res[i][k] = 0;
				}
			}
		}
		return res;
	}

	public MnistDataSet(String path) {
		if (path != null) {
			if (path.charAt(path.length() - 1) == File.separatorChar) {
				this.dataPath = path;
			} else {
				this.dataPath = path + File.separator;
			}
		} else {
			this.dataPath = "." + File.separator;
		}
	}

	public void load_mnist() throws Exception {
		doload(0); // load train data set
		doload(1); // load test data set
	}

	private void doload(int flag) throws Exception {
		String label_File = dataPath + (flag == 0 ? train_labels_idx_file : test_labels_idx_file);
		String image_File = dataPath + (flag == 0 ? train_images_file : test_images_file);
		DataInputStream dinLabel = new DataInputStream(new FileInputStream(label_File)),
				dinImage = new DataInputStream(new FileInputStream(image_File));
		int magicLabelNumber = dinLabel.readInt();
		int magicImageNumber = dinImage.readInt();
		int numLabel = dinLabel.readInt();
		int numImage = dinImage.readInt();
		if (magicLabelNumber != 2049 || magicImageNumber != 2051 || numLabel != numImage) {
			dinLabel.close();
			dinImage.close();
			throw new RuntimeException(dataPath + " doesn't contain a valid mnist data set");
		}

		rows = dinImage.readInt();
		cols = dinImage.readInt();
		int len = rows * cols;
		int[][] image = new int[numImage][len];
		int[] label = new int[numLabel];
		for (int i = 0; i < numImage; i++) {
			label[i] = dinLabel.readUnsignedByte();
			for (int k = 0; k < len; k++) {
				image[i][k] = dinImage.readUnsignedByte();
			}
		}
		dinLabel.close();
		dinImage.close();
		if (flag == 0) {
			train_N = numImage;
			train_Images = image;
			train_Labels = label;
		} else {
			test_N = numImage;
			test_Images = image;
			test_Labels = label;
		}
	}

	public static void main(String[] args) {
		MnistDataSet ds = new MnistDataSet("E:\\ai.projects\\mnist");
		try {
			Date begin = new Date();
			System.out.println("Begin loading...");
			ds.load_mnist();
			System.out.println("Time consumed:" + (new Date().getTime() - begin.getTime()));
			System.out.println("train_N\ttest_N");
			System.out.printf("%d\t%d", ds.getTrain_N(), ds.getTest_N());
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
}
